"""
任务四：迁移诊断可解释性分析主程序（task4-new适配版本）

基于task1-new、task2-new、task3-new的更改，对迁移诊断的可解释性进行全面分析

作者：数学建模团队
版本：3.0 (task4-new适配版本)
"""

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 导入任务四模块
from simplified_interpretability_analysis import (
    SimplifiedInterpretabilityAnalyzer, SimplifiedInterpretabilityVisualizer, 
    generate_interpretability_report
)
from task4_adapted_config import *

# 设置matplotlib非交互式后端
plt.switch_backend('Agg')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 10

def load_adapted_data():
    """加载适配的数据（基于task1-new和task3-new）"""
    print("=" * 60)
    print("步骤1: 加载适配的任务数据")
    print("=" * 60)
    
    # 加载源域数据（来自task1-new）
    source_csv_path = DATA_CONFIG['source_csv_path']
    print(f"📂 加载源域数据: {source_csv_path}")
    
    if not os.path.exists(source_csv_path):
        print(f"❌ 源域数据文件不存在: {source_csv_path}")
        return None
    
    source_df = pd.read_csv(source_csv_path)
    print(f"✅ 源域数据加载成功: {source_df.shape}")
    
    # 加载目标域数据（来自task1-new）
    target_csv_path = DATA_CONFIG['target_csv_path']
    print(f"📂 加载目标域数据: {target_csv_path}")
    
    if not os.path.exists(target_csv_path):
        print(f"❌ 目标域数据文件不存在: {target_csv_path}")
        return None
    
    target_df = pd.read_csv(target_csv_path)
    print(f"✅ 目标域数据加载成功: {target_df.shape}")
    
    # 准备特征数据
    feature_cols = [col for col in source_df.columns if col not in DATA_CONFIG['exclude_columns']]
    print(f"📊 特征列数量: {len(feature_cols)}")
    
    # 源域数据
    X_source = source_df[feature_cols].values
    y_source = source_df['fault_type'].values
    
    # 目标域数据（无标签，需要从task3-new获取预测结果）
    X_target = target_df[feature_cols].values
    
    # 标签编码
    from sklearn.preprocessing import LabelEncoder
    label_encoder = LabelEncoder()
    y_source_encoded = label_encoder.fit_transform(y_source)
    
    # 模拟目标域预测结果（基于task3-new的输出）
    # 这里我们使用一个简单的分类器来生成目标域预测
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.preprocessing import StandardScaler
    
    # 标准化特征
    scaler = StandardScaler()
    X_source_scaled = scaler.fit_transform(X_source)
    X_target_scaled = scaler.transform(X_target)
    
    # 训练简单分类器进行目标域预测
    rf = RandomForestClassifier(n_estimators=100, random_state=42)
    rf.fit(X_source_scaled, y_source_encoded)
    y_target_pred = rf.predict(X_target_scaled)
    
    print(f"📊 数据统计:")
    print(f"   - 源域样本数: {len(X_source)}")
    print(f"   - 目标域样本数: {len(X_target)}")
    print(f"   - 特征维度: {X_source.shape[1]}")
    print(f"   - 类别数量: {len(label_encoder.classes_)}")
    print(f"   - 类别名称: {list(label_encoder.classes_)}")
    print(f"   - 目标域预测分布: {dict(zip(*np.unique(y_target_pred, return_counts=True)))}")
    
    return X_source_scaled, y_source_encoded, X_target_scaled, y_target_pred, label_encoder

def create_adapted_dann_model(input_dim, num_classes, device='cpu'):
    """创建适配的DANN模型（基于task3-new的架构）"""
    print("\n" + "=" * 60)
    print("步骤2: 创建适配的DANN模型")
    print("=" * 60)
    
    # 特征提取器（基于task3-new的简化架构）
    class FeatureExtractor(nn.Module):
        def __init__(self, input_dim, hidden_dims=[64, 32]):
            super(FeatureExtractor, self).__init__()
            layers = []
            prev_dim = input_dim
            
            for hidden_dim in hidden_dims:
                layers.extend([
                    nn.Linear(prev_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Dropout(0.3)
                ])
                prev_dim = hidden_dim
                
            self.feature_extractor = nn.Sequential(*layers)
            self.output_dim = hidden_dims[-1]
            
        def forward(self, x):
            return self.feature_extractor(x)
    
    # 分类器
    class Classifier(nn.Module):
        def __init__(self, input_dim, num_classes):
            super(Classifier, self).__init__()
            self.classifier = nn.Sequential(
                nn.Linear(input_dim, 16),
                nn.ReLU(),
                nn.Dropout(0.3),
                nn.Linear(16, num_classes)
            )
            
        def forward(self, x):
            return self.classifier(x)
    
    # 域判别器
    class DomainDiscriminator(nn.Module):
        def __init__(self, input_dim):
            super(DomainDiscriminator, self).__init__()
            self.discriminator = nn.Sequential(
                nn.Linear(input_dim, 32),
                nn.ReLU(),
                nn.Dropout(0.3),
                nn.Linear(32, 2)
            )
            
        def forward(self, x):
            return self.discriminator(x)
    
    # 梯度反转层
    class GradientReversalLayer(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x, alpha):
            ctx.alpha = alpha
            return x.view_as(x)
        
        @staticmethod
        def backward(ctx, grad_output):
            output = grad_output.neg() * ctx.alpha
            return output, None
    
    # DANN模型
    class DANNModel(nn.Module):
        def __init__(self, input_dim, num_classes):
            super(DANNModel, self).__init__()
            self.feature_extractor = FeatureExtractor(input_dim)
            self.classifier = Classifier(self.feature_extractor.output_dim, num_classes)
            self.domain_discriminator = DomainDiscriminator(self.feature_extractor.output_dim)
            self.gradient_reversal = GradientReversalLayer.apply
            
        def forward(self, x, alpha=1.0):
            features = self.feature_extractor(x)
            class_output = self.classifier(features)
            reversed_features = self.gradient_reversal(features, alpha)
            domain_output = self.domain_discriminator(reversed_features)
            return class_output, domain_output, features
    
    model = DANNModel(input_dim, num_classes).to(device)
    print(f"✅ DANN模型创建成功:")
    print(f"   - 输入维度: {input_dim}")
    print(f"   - 类别数量: {num_classes}")
    print(f"   - 特征维度: {model.feature_extractor.output_dim}")
    print(f"   - 设备: {device}")
    
    return model

def perform_interpretability_analysis(model, X_source, y_source, X_target, y_target_pred, label_encoder, timestamp):
    """执行可解释性分析"""
    print("\n" + "=" * 60)
    print("步骤3: 执行可解释性分析")
    print("=" * 60)
    
    # 创建分析器
    analyzer = SimplifiedInterpretabilityAnalyzer(model, MODEL_CONFIG['device'])
    
    # 1. 特征重要性分析
    print("\n🔍 1. 特征重要性分析")
    feature_importance = analyzer.analyze_feature_importance(X_source, y_source)
    
    # 2. 域适应过程分析
    print("\n🔍 2. 域适应过程分析")
    domain_adaptation = analyzer.analyze_domain_adaptation(X_source, y_source, X_target, y_target_pred)
    
    # 3. 故障特征分析
    print("\n🔍 3. 故障特征分析")
    fault_characteristics = analyzer.analyze_fault_characteristics(X_source, y_source, label_encoder)
    
    # 4. 决策边界分析
    print("\n🔍 4. 决策边界分析")
    decision_boundary = analyzer.analyze_decision_boundary(X_source, y_source, label_encoder)
    
    # 保存分析结果
    analysis_results = {
        'X_source': X_source,
        'y_source': y_source,
        'X_target': X_target,
        'y_target_pred': y_target_pred,
        'feature_importance': feature_importance,
        'domain_adaptation': domain_adaptation,
        'fault_characteristics': fault_characteristics,
        'decision_boundary': decision_boundary
    }
    
    print("✅ 可解释性分析完成")
    return analysis_results

def create_interpretability_visualizations(analysis_results, label_encoder, timestamp):
    """创建可解释性可视化"""
    print("\n" + "=" * 60)
    print("步骤4: 创建可解释性可视化")
    print("=" * 60)
    
    # 创建可视化器
    visualizer = SimplifiedInterpretabilityVisualizer(INTERPRETABILITY_CONFIG)
    
    # 1. 特征重要性可视化
    print("\n📊 1. 特征重要性可视化")
    visualizer.plot_feature_importance(analysis_results['feature_importance'], timestamp)
    
    # 2. 域适应过程可视化
    print("\n📊 2. 域适应过程可视化")
    visualizer.plot_domain_adaptation(analysis_results['domain_adaptation'], timestamp)
    
    # 3. 故障特征可视化
    print("\n📊 3. 故障特征可视化")
    visualizer.plot_fault_characteristics(analysis_results['fault_characteristics'], timestamp)
    
    # 4. 决策边界可视化
    print("\n📊 4. 决策边界可视化")
    visualizer.plot_decision_boundary(analysis_results['decision_boundary'], timestamp)
    
    print("✅ 可解释性可视化完成")

def generate_adapted_report(analysis_results, label_encoder, timestamp):
    """生成适配报告"""
    print("\n" + "=" * 60)
    print("步骤5: 生成适配报告")
    print("=" * 60)
    
    # 调用报告生成函数
    generate_interpretability_report(analysis_results, label_encoder, timestamp)

def main():
    """主函数"""
    print("=" * 80)
    print("🚀 任务四：迁移诊断可解释性分析（task4-new适配版本）")
    print("=" * 80)
    
    # 生成时间戳
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    try:
        # 1. 加载适配的数据
        data_result = load_adapted_data()
        if data_result is None:
            print("❌ 数据加载失败，程序退出")
            return
        
        X_source, y_source, X_target, y_target_pred, label_encoder = data_result
        
        # 2. 创建适配的DANN模型
        model = create_adapted_dann_model(
            MODEL_CONFIG['input_dim'], 
            MODEL_CONFIG['num_classes'], 
            MODEL_CONFIG['device']
        )
        
        # 3. 执行可解释性分析
        analysis_results = perform_interpretability_analysis(
            model, X_source, y_source, X_target, y_target_pred, label_encoder, timestamp
        )
        
        # 4. 创建可解释性可视化
        create_interpretability_visualizations(analysis_results, label_encoder, timestamp)
        
        # 5. 生成适配报告
        generate_adapted_report(analysis_results, label_encoder, timestamp)
        
        print("\n" + "=" * 80)
        print("🎉 任务四：迁移诊断可解释性分析（task4-new适配版本）完成！")
        print("=" * 80)
        print(f"📁 生成文件:")
        print(f"   - feature_importance_analysis_{timestamp}.png")
        print(f"   - domain_adaptation_analysis_{timestamp}.png")
        print(f"   - fault_characteristics_analysis_{timestamp}.png")
        print(f"   - decision_boundary_analysis_{timestamp}.png")
        print(f"   - task4_adapted_interpretability_report_{timestamp}.md")
        
    except Exception as e:
        print(f"❌ 程序执行出错: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
